import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

from planet import *


fig = plt.figure(figsize=[6, 6])
ax = fig.add_subplot(111, projection='3d')

color_palette = ['blue', 'green', 'red', 'orange']

class AnimManager:
    '''
    补间
    '''

    def __init__(self):
        self. plantManager = PlanetManager()
        self.lines = []

        for i, planet in enumerate(self.plantManager.planets):
            poss = planet.trajectory
            plot, = ax.plot(poss[:, 0], poss[:, 1], poss[:, 2], color=color_palette[i])
            self.lines.append(plot)

    def update(self):
        '''
        迭代 24*30*4 次
        '''
        # for _ in range(24*30*4):
        self.plantManager.update_state()

    def update_lim(self):
        # 存储所有行星的坐标到数组中
        x = []
        y = [] 
        z = []
        for planet in self.plantManager.planets:
            x += planet.trajectory[:,0].tolist() 

            y += planet.trajectory[:,1].tolist()
            z += planet.trajectory[:,2].tolist()
            
        # 设置坐标轴范围 
        xmin, xmax = np.min(x), np.max(x)
        ymin, ymax = np.min(y), np.max(y) 
        zmin, zmax = np.min(z), np.max(z)
        margin = 0.05 # 保证一定空白边界

        ax.set_xlim(xmin-margin, xmax+margin)  
        ax.set_ylim(ymin-margin, ymax+margin)
        ax.set_zlim(zmin-margin, zmax+margin)
        
    def animate(self, frame_num):
        am.update()
        am.update_lim()
        for i, planet in enumerate(self.plantManager.planets):
            plot = self.lines[i]
            poss = planet.trajectory
            plot.set_data(poss[:frame_num, 0], poss[:frame_num, 1])
            plot.set_3d_properties(poss[:frame_num, 2])
            # plot.scatter([0], [0], c='r', marker='*')
        
        return self.lines

am = AnimManager()


# ax.set_xlim(-1e10, 1e10)  
# ax.set_ylim(-1e10, 1e10)  
# ax.set_zlim(-1e10, 1e10)  
# am.update()
# am.update_lim()
ani = animation.FuncAnimation(fig, am.animate, frames=None, interval=40, save_count= 100)
# ani = animation.FuncAnimation(fig, am.animate, frames=24*30*4, interval=1, blit=True)

plt.show()
        




